from typing import Callable
from learners.multi_base_ca import Learner as BaselineLearner

# from learners.baseline import ClsAlign
from models.multi_net import MultiRoute

num_workers = 8


class Learner(
    BaselineLearner[MultiRoute],
):
    _network: MultiRoute

    def __init__(self, args, data_manager, model_func: Callable = MultiRoute):
        super().__init__(args, data_manager, model_func)
        self.evalNME = False
        self.ca_epochs = args["ca_epochs"]
        self.ca_lr = args["ca_lr"]
        self.ca_forward = lambda model, inputs: model(
            inputs, mode="cur", ca=True
        )

    def _update_fc(self, nb_classes):
        self._network.update_fc(
            nb_classes,
            freeze_old=True,
            fc_kwargs={
                "init_experts": self.args["init_experts"],
                "fc_temperture": self.args["fc_temperture"],
                "with_norm": self.args["with_norm"],
            },
        )
